from socket import *
from threading import Thread
import threading
import rtmp_msg_config as rmc
import rtmp_streamm as rst


# rtmp head 参数格式:
# chunk stream 1个字节(\x02)
# timestamp 时间戳 3个字节 (\x00\x00\x00)
# body size body请求体的大小 3个字节(\x00\x00\x04) 这里设置请求体几个字节 下面请求体就要是几个字节
# type id 类型 windows ack = \x05, set peer bandwidth = \x06 , set chunk size= \x01,
# amf0 command _result('NetConnection.Connect.Success') = \x14
# Stream ID通常用以完成某些特定的工作，如使用ID为0的Stream来完成客户端和服务器的连接和控制，
# 使用ID为1的Stream来完成视频流的控制和播放等工, stream为 4个字节

host = ''
port = 1935
ADDR = (host, port)
tctime = socket(AF_INET, SOCK_STREAM)
tctime.bind(ADDR)
tctime.listen(3)

nums = 1
mutex = threading.Lock()

video_list = []
start_video = []
pushonmetadata = {}
rtmp_client_queue = []
threading_client_video_list = []


# 握手连接
def hand_shake(rtmp_client):
    # 接收 c0c1 握手数据
    data = rtmp_client.recv(1537)

    # 如果默认用 bytes(data[0]) 的话，python 会将int 转 二进制转为 \x00 所有要用int类型中的to_bytes方法
    s0 = data[0].to_bytes(1, "big")

    s1 = b""
    for i in range(0, 1536):  # 握手数据长度是固定的 1536
        s1 += b'\x00'

    # 握手 服务器发送的 s0s1s2 数据格式是 c0c1[0] + 数据生成的1536个随机二进制数据 + c0c1
    send_data = s0 + s1 + data

    # 发送 s0s1s2
    rtmp_client.send(send_data)

    # 接收c2数据
    s2 = rtmp_client.recv(len(send_data))


# 接收chuck数据
def chuck_connect(rtmp_client):
    data = b""
    while 1:
        receive = rtmp_client.recv(1024 * 10)
        data += receive
        print("chuck_connect", receive)
        print("data,   ", data)
        if len(receive) == 0:
            break
        if rmc.ConnectEnd in receive: break
    if len(data) == 0:
        return "error"
    if b"video" in data:
        return "pull"
    return "push"


# 创建流
def create_stream(rtmp_client):
    while 1:
        receive = rtmp_client.recv(1024)
        # 接收创建流请求，
        if rmc.CreateStream in str(receive):
            break


# 构造响应数据
def net_connection_result():
    response = {"fmsVer": "FMS/3,0,1,123", "capabilities": float(31)}
    event = {"level": "status", "code": "NetConnection.Connect.Success",
             "description": "Connection succeeded.",
             "objectEncoding": 0}

    return rst.NetConnectionResult(response, event)


def pulldataframe(params):
    global pushonmetadata
    pushonmetadata = rst.PullDataFrame(params)


def create_steam(pl):
    result = ["_result", 2 if pl == "pull" else 4, None]
    return rst.CreateSteam(result)


def publish_receive(rtmp_client):
    while 1:
        receive = rtmp_client.recv(1024)
        print("publish", receive)
        if rmc.Publish in str(receive):
            break


# 推流开始
def netstream_publish_start():
    data = ["onStatus", 0, None,
            {"description": "Start publishing", "level": "status", "code": "NetStream.Publish.Start"}]
    return rst.OnStatus(data)


def receive_rtmp_data(rtmp_client):
    data = b""
    global video_data, video_list, pushonmetadata, threading_client_video_list

    while 1:
        receive = rtmp_client.recv(1024 * 50)
        if not data:
            data = receive
            pushonmetadata = rst.PullDataFrame(receive)
        else:
            if len(start_video) < 2:
                start_video.append(receive)
            else:
                mutex.acquire()
                tcvl = threading_client_video_list[:]
                mutex.release()
                for i in tcvl:
                    i.append(receive)

        video_data = receive
        if len(receive) == 0:
            break


def playreset():
    data = ["onStatus", 0, None,
            {"level": "status", "code": "NetStream.Play.Reset", "description": "Playing and resetting stream"}]
    return rst.PlayReset(data)


def playstart():
    data = ["onStatus", 0, None,
            {"level": "status", "code": "NetStream.Play.Start", "description": "Started playing stream"}]
    return rst.PlayReset(data)


def datastart():
    data = ["onStatus", 0, None,
            {"level": "status", "code": "NetStream.Data.Start", "description": "Started playing stream"}]
    return rst.PlayReset(data)


def playpublishnotify():
    data = ["onStatus", 0, None,
            {"level": "status", "code": "NetStream.Play.PublishNotify", "description": "Started playing notify"}]
    return rst.PlayReset(data)


def onmetadata():
    data = ["onMetaData", [{"duration": float(0)}, {"fileSize": float(0)}, {"width": float(1280)},
                           {"height": float(720)}, {"videocodecid": "avc1"}, {"videodatarate": float(2500)},
                           {"framerate": float(30)}, {"audiocodecid": "mp4a"}, {"audiodatarate": float(160)},
                           {"audiosamplerate": float(44100)}, {"audiosamplesize": float(16)},
                           {"audiochannels": float(2)}, {"stereo": True}, {"2.1": False}, {"3.1": False},
                           {"4.0": False}, {"4.1": False}, {"5.1": False}, {"7.1": False},
                           {"encoder": "obs-output module (libobs version 25.0.8)"}]]

    return rst.OnMetaData(data)


def onStatus():
    data = ["onStatus", 0, None, {"level": "status", "code": "NetStream.Play.Start", "description": "Start live"}]
    return rst.OnStatus(data)


def RtmpSampleAccess():
    data = ["|RtmpSampleAccess", True, True]
    return rst.RtmpSampleAccess(data)


def StreamBegin():
    data = []
    return rst.StreamBegin(data)


def OnMetaData():
    data = ["onMetaData", {"Server": "Python nanshen yyds", "width": pushonmetadata.get("width", 1280.0),
                           "height": pushonmetadata.get("height", 720.0),
                           "displayWidth": pushonmetadata.get("width", 1280.0),
                           "displayHeight": pushonmetadata.get("height", 720.0),
                           "duration": pushonmetadata.get("duration", 0.0),
                           "framerate": pushonmetadata.get("framerate", 30.0),
                           "fps": pushonmetadata.get("framerate", 30.0),
                           "videodatarate": pushonmetadata.get("videodatarate", 5000.0),
                           "videocodecid": pushonmetadata.get("videocodecid", 7.0),
                           "audiodatarate": pushonmetadata.get("audiodatarate", 160.0),
                           "audiocodecid": pushonmetadata.get("audiocodecid", 10.0),
                           "profile": pushonmetadata.get("profile", ""), "level": pushonmetadata.get("level", ""),
                           }]
    return rst.OnMetaData(data)


def pull(rtmp_client):
    while 1:
        receive = rtmp_client.recv(1024 * 10)
        if rmc.PlayCreeateStream in receive and b'createStream' in receive: break

    # 响应创建流
    send_data = create_steam("pull")
    rtmp_client.send(send_data)

    # 接收 play 拉流请求
    while 1:
        receive = rtmp_client.recv(1024 * 10)
        if rmc.PlayEnd in receive:
            break

    rtmp_client.send(onStatus())

    rtmp_client.send(RtmpSampleAccess())

    rtmp_client.send(StreamBegin())
    aa = OnMetaData()
    rtmp_client.send(aa)
    global threading_client_video_list

    mutex.acquire()
    video_list_index = len(threading_client_video_list)
    threading_client_video_list.append([])
    mutex.release()

    for i in start_video:
        rtmp_client.send(i)

    global video_list
    while 1:
        send_data_list = threading_client_video_list[video_list_index][:]
        threading_client_video_list[video_list_index] = []
        index = 0
        for i in send_data_list:
            print("start", len(send_data_list))
            try:
                rtmp_client.send(i)
                print("hhhhhhhh", i)
            except:
                threading_client_video_list.pop(video_list_index)
                rtmp_client.close()
                print("发送错误", i)
                print(send_data_list[index-1])
                return
            print("end")
            index += 1


def client_push():
    global video_list
    cc_list = []

    while 1:
        mutex.acquire()
        send_data_list = video_list[:]
        video_list = []
        mutex.release()
        index = 0
        for i in send_data_list:
            if i not in cc_list:
                mutex.acquire()
                ls_threading_queue = rtmp_client_queue[:]
                mutex.release()
                ls_thread_list = []
                for rtmp_client in ls_threading_queue:
                    t = threading.Thread(target=rtmp_client_pop, args=(rtmp_client, i))
                    t.start()
                    ls_thread_list.append(t)
                for rtmp_client_join in ls_thread_list:
                    print("join")
                    rtmp_client_join.join()

            index += 1
        cc_list = send_data_list[:]


def rtmp_client_pop(rtmp_client, send_data):
    global rtmp_client_queue
    if rtmp_client not in rtmp_client_queue:
        return
    rtmp_client_tcp = ""
    for i in rtmp_client.keys():
        rtmp_client_tcp = i
    try:
        rtmp_client_tcp.send(send_data)
    except error:
        rtmp_client[rtmp_client_tcp] += 1
        if rtmp_client[rtmp_client_tcp] >= 5:
            mutex.acquire()
            rtmp_client_queue.remove(rtmp_client)
            rtmp_client_tcp.close()
            mutex.release()


start = 0


def cl_data(rtmp_client):
    # 握手连接
    hand_shake(rtmp_client)

    # 接收chuck 和 connect数据
    pl = chuck_connect(rtmp_client)
    if pl == "error":
        rtmp_client.close()
        return

    # Window Acknowledgement  Size消息  确认窗口大小
    window_ack = rst.SetWindowAck(5000000)

    bb = rtmp_client.send(window_ack)

    # 设置绑定窗口大小
    set_peer_act = rst.SetPeer(5000000)
    # rtmp_client.send(set_peer_act)

    # # 流媒体数据块大小
    set_chunk_size = rst.SetChunkSzie(4096)
    # rtmp_client.send(set_chunk_size)

    # 响应连接成功 net_connection_result 数据
    net_result_msg = net_connection_result()
    # rtmp_client.send(net_result_msg)

    # # 设置窗口大小，流大小， 创建块成功
    send_data = set_peer_act + set_chunk_size + net_result_msg
    bb = rtmp_client.send(send_data)

    if pl == "pull":
        pull(rtmp_client)
    else:
        # rtmp_client.send(send_data)
        # 接收创建流请求
        create_stream(rtmp_client)
        send_data = create_steam("push")
        rtmp_client.send(send_data)
        # 接收创建流请求
        publish_receive(rtmp_client)
        send_data = netstream_publish_start()
        rtmp_client.send(send_data)

        receive_rtmp_data(rtmp_client)
        rtmp_client.close()


while True:
    tctimeClient, addr = tctime.accept()
    t1 = Thread(target=cl_data, args=(tctimeClient,))
    t1.start()
